import sys
import os
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QPushButton, QLabel, QFileDialog, QMessageBox
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt
import timm
import torch.nn as nn

# 加载模型
def load_model(model_path, device, n_class):
    # 加载模型权重
    checkpoint = torch.load(model_path, map_location=device, weights_only=True)
    
    # 创建模型架构
    model = timm.create_model('resnest50d', pretrained=False)
    num_features = model.fc.in_features
    
    # 替换全连接层以匹配保存的权重
    model.fc = nn.Linear(num_features, n_class)  # n_class 必须与保存权重时的类别数量一致
    model.load_state_dict(checkpoint['model_state_dict'])  # 加载权重
    model.to(device)
    model.eval()
    return model

# 图像预处理
def preprocess_image(image_path):
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)  # 添加批次维度
    return image

# 图像识别
def recognize_image(model, image_tensor, device, idx_to_labels):
    with torch.no_grad():
        image_tensor = image_tensor.to(device)
        outputs = model(image_tensor)
        _, predicted = torch.max(outputs, 1)
        predicted_label = idx_to_labels[predicted.item()]
    return predicted_label

# PyQt5 程序
class ImageRecognitionApp(QWidget):
    def __init__(self, model_path, idx_to_labels_path, device, n_class):
        super().__init__()
        self.model_path = model_path
        self.idx_to_labels_path = idx_to_labels_path
        self.device = device
        self.n_class = n_class
        self.initUI()
        self.load_model_and_labels()

    def initUI(self):
        self.setWindowTitle('图像识别')
        self.setGeometry(100, 100, 400, 300)

        layout = QVBoxLayout()

        self.image_label = QLabel(self)
        self.image_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.image_label)

        self.result_label = QLabel(self)
        self.result_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.result_label)

        self.upload_button = QPushButton('上传图片', self)
        self.upload_button.clicked.connect(self.upload_image)
        layout.addWidget(self.upload_button)

        self.recognize_button = QPushButton('识别图片', self)
        self.recognize_button.clicked.connect(self.recognize)
        layout.addWidget(self.recognize_button)

        self.setLayout(layout)

    def load_model_and_labels(self):
        self.model = load_model(self.model_path, self.device, self.n_class)
        self.idx_to_labels = np.load(self.idx_to_labels_path, allow_pickle=True).item()

    def upload_image(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.jpg *.jpeg *.bmp);;All Files (*)", options=options)
        if file_path:
            self.image_path = file_path
            pixmap = QPixmap(file_path)
            self.image_label.setPixmap(pixmap.scaled(300, 300, Qt.KeepAspectRatio))
            self.result_label.setText("")

    def recognize(self):
        if not hasattr(self, 'image_path'):
            QMessageBox.warning(self, "警告", "请先上传一张图片！")
            return

        image_tensor = preprocess_image(self.image_path)
        predicted_label = recognize_image(self.model, image_tensor, self.device, self.idx_to_labels)
        self.result_label.setText(f"识别结果: {predicted_label}")

if __name__ == '__main__':
    app = QApplication(sys.argv)

    # 模型路径和类别映射路径
    model_path = 'best50_model.pth'  # 替换为您的模型路径
    idx_to_labels_path = 'idx_to_labels.npy'  # 替换为您的类别映射路径
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    
    # 确保 n_class 与保存权重时的类别数量一致
    n_class = 1781  # 根据保存的权重修改为正确的类别数量

    window = ImageRecognitionApp(model_path, idx_to_labels_path, device, n_class)
    window.show()
    sys.exit(app.exec_())